import itertools
from itertools import islice
import heapq

from algorithm.list import find_if
from number_theory.gcd import gcd
from number_theory.nth_root import int_sqrt
from number_theory.nth_root import int_sqrt_if_perfect
from number_theory.nth_root import is_perfect_square
from prime.prime import is_prime
from prime.prime import Prime

__all__ = ["factor", "pretty_print"]



Prime.init_primes(10000)


def factor(n):
    if n == 0:
        return [(0, 1)]
    elif n < 0:
        return [(-1, 1)] + factor(-n)
    elif n == 1:
        return []
    elif is_prime(n):
        return [(n, 1)]
    
    p = trial(n)
    if p:
        k = 0
        while True:
            b, r = divmod(n, p)
            if r:
                break
            else:
                k += 1
                n = b
        return [(p, k)] + factor(n)

    m1 = factor_big_numbers(n)
    m2 = n // m1
    return list(heapq.merge(factor(m1), factor(m2)))


def trial(n):
    return find_if(lambda p: n % p == 0, Prime.small_primes)


def factor_big_numbers(n):
    n1 = pollard_rho(n, 0x1000)
    if n1:
        return n1
    for m in (1, 3, 5, 7, 11, 15, 21, 33, 35, 55, 77,
            105, 165, 231, 385, 1155):
        q = squfof(n, m)
        if q:
            return q
    return squfof_old(n)


def pollard_rho(n, limit = None):
    def initial():
        x0 = 0
        a = 0
        while True:
            if x0 == a:
                x0 += 1
                a = 1
            else:
                a += 1
            yield x0, a, lambda x: (x ** 2 + a) % n

    def exp2():
        k = 1
        while limit is None or k < limit:
            yield k
            k <<= 1

    for x0, a, f in initial():
        x = f(x0)
        y = f(x)
        g = gcd(x - y, n)
        if g == n:
            continue
        elif g > 1:
            return g
        for k in exp2():
            fail = False
            x = y
            for i in range(k):
                y = f(y)
            for i in range(k):
                y = f(y)
                g = gcd(x - y, n)
                if g == n:
                    fail = True
                    break
                elif g > 1:
                    return g
            if fail:
                break
        else:
            break

    return None


def squfof_old(n):
    sqrt = int_sqrt_if_perfect(n)
    if sqrt:
        return sqrt

    for N in itertools.count(n, n):
        m = int_sqrt(N)
        m1 = 0
        p0 = 0
        p1 = m
        q0 = 1
        q1 = N - m ** 2
        while True:
            p0, p1 = p1, p0
            q0, q1 = q1, q0
            b, r = divmod(m + p0, q0)
            p1 = m - r
            q1 += b * (p0 - p1)
            m1 = int_sqrt_if_perfect(q1)
            if m1:
                break
        q0 = m1
        p1 = m - (m - p1) % q0
        q1 = (N - p1 ** 2) // q0
        while True:
            p0, p1 = p1, p0
            q0, q1 = q1, q0
            b, r = divmod(m + p0, q0)
            p1 = m - r
            if p0 == p1:
                break
            q1 += b * (p0 - p1)
        g = gcd(n, p0)
        if 1 < g < n:
            return g


def squfof(n, m = 1):
    sqrt = int_sqrt(n)
    if sqrt ** 2 == n:
        return sqrt

    d = m * n
    if d & 3 == 1:
        d <<= 1
    s = int_sqrt(d)
    q_hat = 1
    p = s
    q = d - p ** 2
    l = int_sqrt(int_sqrt(d) << 1) << 1
    b = l << 1
    queue = []

    for i in range(b + 1):
        little_q, r0 = divmod(s + p, q)
        p_prime = s - r0
        g = q // gcd(q, m << 1)
        if g <= l:
            queue.append((g, p % g))
        q, q_hat = q_hat + little_q * (p - p_prime), q
        p = p_prime
        if i & 1 == 0:
            r = int_sqrt_if_perfect(q)
            if r is not None:
                for z, (x, y) in enumerate(queue):
                    if x == r and (p - y) % r == 0:
                        break
                else:
                    break
                if r > 1:
                    del queue[:z + 1]
                elif r == 1:
                    return None
    else:
        return None;

    q_hat = int_sqrt(q)
    p = s - (s - p) % q_hat
    q = (d - p ** 2) // q_hat

    for i in range(1, b):
        little_q, r0 = divmod(s + p, q)
        p_prime = s - r0
        if p == p_prime:
            break
        q, q_hat = q_hat + little_q * (p - p_prime), q
        p = p_prime

    return q // gcd(q, m << 1)


def pretty_print(factor_result):
    if not factor_result:
        return "1"
    l = []
    for p, k in factor_result:
        if k == 1:
            l.append(str(p))
        else:
            l.append("{0}^{1}".format(p, k))
    return " * ".join(l)
